from __future__ import annotations

from dataclasses import dataclass
from typing import Iterable, Literal, Sequence

import numpy as np


StateLayout = Literal["row", "column"]
ControlLayout = Literal["row", "column"]


@dataclass
class SimulationContext:
    t0: float
    x0: np.ndarray
    x_initial: np.ndarray
    x_guess: np.ndarray
    next_states: np.ndarray
    xs: np.ndarray
    u0: np.ndarray
    state_history: list
    control_history: list
    time_history: list
    trajectory_log: list
    sim_time: float


def _prepare_state_vector(
    values: Sequence[float] | None, size: int, *, fallback: Iterable[float]
) -> np.ndarray:
    if values is None:
        values = list(fallback)
    array = np.asarray(values, dtype=float).reshape(-1)
    if array.size != size:
        raise ValueError(f"Expected {size} state entries, got {array.size}")
    return array


def _prepare_control_seed(seed: Sequence[float] | None, size: int) -> np.ndarray:
    if seed is None:
        seed = [0.0] * size
    array = np.asarray(seed, dtype=float).reshape(-1)
    if array.size != size:
        raise ValueError(f"Expected {size} control entries, got {array.size}")
    return array


def create_simulation_context(
    *,
    n_states: int,
    n_controls: int,
    N: int,
    initial_state: Sequence[float] | None = None,
    final_state: Sequence[float] | None = None,
    control_seed: Sequence[float] | None = None,
    state_layout: StateLayout = "column",
    control_layout: ControlLayout = "row",
    sim_time: float = 20.0,
    include_initial_time: bool = False,
    default_final_state: Iterable[float] = (1.5, 1.5, 0.0),
) -> SimulationContext:
    """
    Provide reusable simulation scaffolding (initial guesses, logs, and bounds).
    """
    t0 = 0.0
    x0 = _prepare_state_vector(initial_state, n_states, fallback=[0.0] * n_states)
    xs = _prepare_state_vector(
        final_state,
        n_states,
        fallback=list(default_final_state)[:n_states]
        + [0.0] * max(0, n_states - len(default_final_state)),
    )
    control_base = _prepare_control_seed(control_seed, n_controls)

    if state_layout == "row":
        x_guess = np.zeros((N + 1, n_states))
    elif state_layout == "column":
        x_guess = np.zeros((n_states, N + 1))
    else:
        raise ValueError(f"Unsupported state_layout '{state_layout}'")
    next_states = np.zeros_like(x_guess)

    if control_layout == "row":
        u0 = np.tile(control_base, N).reshape(N, n_controls)
    elif control_layout == "column":
        u0 = np.tile(control_base, N).reshape(n_controls, N)
    else:
        raise ValueError(f"Unsupported control_layout '{control_layout}'")

    time_history = [t0] if include_initial_time else []

    return SimulationContext(
        t0=t0,
        x0=x0,
        x_initial=x0.copy(),
        x_guess=x_guess,
        next_states=next_states,
        xs=xs,
        u0=u0,
        state_history=[],
        control_history=[],
        time_history=time_history,
        trajectory_log=[],
        sim_time=sim_time,
    )
